// ----------------------------------------------------------------------------
// diabindLoader.h
// Load debug information
// ----------------------------------------------------------------------------

#pragma once

// ----------------------------------------------------------------------------

#include "diabindSession.h"
#include "diabindEnumSymbols.h"

// ----------------------------------------------------------------------------

namespace diabind
{

// ----------------------------------------------------------------------------

class CLoader
{
public:
    // Load debug information from PDB-file
    static CSessionPtr FromPdb(const std::wstring &wstrFilePath);
    static CSessionPtr FromExe(const std::wstring &wstrFilePath, const std::wstring &wstrSearchPath);

protected:
    // Exception: DIA interface not registered
    class CDiaSourceNotRegistered{};

    static CSessionPtr FromPdb80(const std::wstring &wstrFilePath);
    static CSessionPtr FromExe80(const std::wstring &wstrFilePath, const std::wstring &wstrSearchPath);
    static CSessionPtr FromPdb90(const std::wstring &wstrFilePath);
    static CSessionPtr FromExe90(const std::wstring &wstrFilePath, const std::wstring &wstrSearchPath);
    static CSessionPtr FromPdb100(const std::wstring &wstrFilePath);
    static CSessionPtr FromExe100(const std::wstring &wstrFilePath, const std::wstring &wstrSearchPath);
    static CSessionPtr FromPdb110(const std::wstring &wstrFilePath);
    static CSessionPtr FromExe110(const std::wstring &wstrFilePath, const std::wstring &wstrSearchPath);
    static CSessionPtr FromPdb120(const std::wstring &wstrFilePath);
    static CSessionPtr FromExe120(const std::wstring &wstrFilePath, const std::wstring &wstrSearchPath);

    template< class DiaDataSourceT, class DiaSessionT >
    static CSessionPtr OpenGlobalSession(
        DiaDataSourceT *pDataSource,
        HRESULT (STDMETHODCALLTYPE DiaDataSourceT::*pMethod)(DiaSessionT **)
    )
    {
        _com_ptr_t< _com_IIID< DiaSessionT, &__uuidof(DiaSessionT) > > pSession;
        HRESULT hResult = (pDataSource->*pMethod)(&pSession);
        if (S_OK != hResult)
            throw CException(hResult, "IDiaDataSource::openSession");

        return CSessionPtr( new CSessionImpl< DiaSessionT >(pSession) );
    }

    template < class DiaDataSourceT >
    class CLoaderFromPdb : boost::noncopyable
    {
    public:
        CLoaderFromPdb(const std::wstring &wstrFilePath) 
            : m_wstrFilePath(wstrFilePath)
        {
        }

        HRESULT operator()(DiaDataSourceT *pDiaDataSource)
        {
            return pDiaDataSource->loadDataFromPdb(m_wstrFilePath.c_str());
        }

    private:
        const std::wstring &m_wstrFilePath;
    };

    template< class DiaLoadCallback2T, class DiaLoadCallbackT >
    class CEmptyCallback : public DiaLoadCallback2T
    {
    public:
        ULONG STDMETHODCALLTYPE AddRef() { return 2; }
        ULONG STDMETHODCALLTYPE Release() { return 1; }
        HRESULT STDMETHODCALLTYPE QueryInterface( REFIID rid, void **ppUnk )
        {
            if (!ppUnk)
                return E_INVALIDARG;
            *ppUnk = NULL;
            if (rid == __uuidof(DiaLoadCallback2T))
                *ppUnk = reinterpret_cast<DiaLoadCallback2T *>(this);
            else if (rid == __uuidof(DiaLoadCallbackT))
                *ppUnk = reinterpret_cast<DiaLoadCallbackT *>(this);
            else if (rid == __uuidof(IUnknown) )
                *ppUnk = reinterpret_cast<IUnknown *>(this);
            return (*ppUnk) ? S_OK : E_NOINTERFACE;
        }
        HRESULT STDMETHODCALLTYPE NotifyDebugDir(BOOL, DWORD, BYTE []) { return S_OK; }
        HRESULT STDMETHODCALLTYPE NotifyOpenDBG(LPCOLESTR, HRESULT)  { return S_OK; }
        HRESULT STDMETHODCALLTYPE NotifyOpenPDB(LPCOLESTR, HRESULT) { return S_OK; }
        HRESULT STDMETHODCALLTYPE RestrictRegistryAccess() { return S_OK; }
        HRESULT STDMETHODCALLTYPE RestrictSymbolServerAccess() { return S_OK; }
        HRESULT STDMETHODCALLTYPE RestrictOriginalPathAccess() { return S_OK; }
        HRESULT STDMETHODCALLTYPE RestrictReferencePathAccess() { return S_OK; }
        HRESULT STDMETHODCALLTYPE RestrictDBGAccess() { return S_OK; }
        HRESULT STDMETHODCALLTYPE RestrictSystemRootAccess() { return S_OK; }
    };

    template < class DiaDataSourceT, class DiaLoadCallback2T, class DiaLoadCallbackT >
    class CLoaderFromExe : boost::noncopyable
    {
    public:
        CLoaderFromExe(const std::wstring &wstrFilePath, const std::wstring &wstrSearchPath) 
            : m_wstrFilePath(wstrFilePath)
            , m_wstrSearchPath(wstrSearchPath)
        {
        }

        HRESULT operator()(DiaDataSourceT *pDiaDataSource)
        {
            CEmptyCallback<DiaLoadCallback2T, DiaLoadCallbackT> EmptyCallback;
            return 
                pDiaDataSource->loadDataForExe(
                    m_wstrFilePath.c_str(),
                    m_wstrSearchPath.c_str(),
                    &EmptyCallback);
        }

    private:
        const std::wstring &m_wstrFilePath;
        const std::wstring &m_wstrSearchPath;
    };

    template< class DiaDataSourceT, class DiaSourceT, class LoaderT >
    static CSessionPtr FromT(LoaderT &Loader, const char *szMethodName)
    {
        _com_ptr_t< _com_IIID< DiaDataSourceT, &__uuidof(DiaDataSourceT) > > pDataSource;
        HRESULT hResult = 
            pDataSource.CreateInstance(
                __uuidof(DiaSourceT),
                NULL,
                CLSCTX_INPROC_SERVER);

        if (REGDB_E_CLASSNOTREG == hResult)
            throw CDiaSourceNotRegistered();

        if (S_OK != hResult)
            throw CException(hResult, "CoCreateInstance");

        hResult = Loader(pDataSource);
        if (S_OK != hResult)
            throw CException(hResult, szMethodName);

        return 
            OpenGlobalSession(
                pDataSource.GetInterfacePtr(),
                &DiaDataSourceT::openSession);
    }
};

// ----------------------------------------------------------------------------

}   // namespace diabind

// ----------------------------------------------------------------------------
